#!/bin/bash

# 检查是否提供了分支名作为参数
if [ -z "$1" ]; then
  echo "Usage: $0 <branch_name>"
  exit 1
fi

BRANCH_NAME=$1

# 更新并初始化子模块
git submodule update --init

# 切换到 Megatron 子模块并检出指定分支
cd Megatron || { echo "Megatron submodule not found!"; exit 1; }
git checkout $BRANCH_NAME || { echo "Failed to checkout branch $BRANCH_NAME in Megatron!"; exit 1; }
cd ..

# 切换到 Megatron-Extension 子模块并检出指定分支
cd Megatron-Extension || { echo "Megatron-Extension submodule not found!"; exit 1; }
git checkout master || { echo "Failed to checkout branch $BRANCH_NAME in Megatron-Extension!"; exit 1; }
cd ..

echo "Submodules have been updated and checked out to branch $BRANCH_NAME."
